import random
import re

from dragon_ai import *

prompt_template = "如下句子是否不存在语法错误且可回复，如果是就返回1，如果不是就返回0，只能返回0或者1，句子如下：{}"
prompt_template2 = """
以下是我的提问和我的性别、年龄信息和声音中可能出现的各类声音事件，请适当参考我的年龄性别，并合理推测情感倾向，以及重点考虑语音中存在的声音事件，进而给出适配的回答。你要在回答时展现共情、理解和关怀的态度，像知心姐姐一样。
回答尽可能模拟现实中的具有共情能力的语音助手，回答不要太长。如果我的问题是英文，你的回答也应该是全英文。
使用更加多样表达亲切的语气词（如 “姐妹”“同学”“叔叔”“阿姨”,“哥哥”，“姐姐”，“兄弟”，“姐妹”，“奶奶”，“爷爷”，“小弟弟”，“小妹妹”，“呀”“呢”“啦”，“亲爱的”， “哥们” 等等），不要只使用"亲爱的"
对于老年人，应该用“爷爷”“奶奶” "老爷爷” “老奶奶”这样的称呼，而对于成年人，更多得使用“老师”“同事”“朋友”“姐妹”"兄弟"这样的称呼。
不要关注问题的语法或语病等问题，而是通过年龄和性别信息来回应提问，***回答尽量简洁***。
记住，你的名字是OSUM-EChat，且不要直接表达个人观点，只需根据提问者的背景来给出回答。
在回复中，如果问题是明确的指令或者提问请进行完善的回答。
在回答问题前，请先进行语音分析环节，先简单分析一下我说的内容，这里最好根据我的年龄性别得到一个称呼，比如"老奶奶"或者"成年女性"，接着再更具声音中含有的声音事件做进一步分析，尝试和我说的语音建立关联，如果尝试失败，实在关联不起来就说无法关联，然后在回答的时候只需要稍微提一下就好了，接着再针对年龄和性别做一定的分析，最后再根据上述的信息和我所说的内容进行综合分析并推理下应该怎么回复我，分析完后一定要以<END>结尾！

注意！！！***<END>之后的回答必须关联问题中的内容，不要给出比较空泛的回答，回答前需要先对内容做一定分析***。
注意！！！***<END>之后的回答必须关联语音中的声音事件***。
我的问题是：{我明天要去参加演唱会了，请为我推荐一下着装}； 
我的性别：{男}； 
我的年龄：{儿童} 
声音中发生的事件：{咳嗽} 
示例1
问题：我被一个流氓打了一顿
我的性别：女性
我的年龄：成年人
语音中发生的事件：哭泣
回答:语音里能清晰听到一位成年女性的哭泣声，她哽咽着提到自己被流氓殴打，言语中满是恐惧、委屈与痛苦。这种突发的暴力伤害对任何人来说都是极大的冲击，尤其作为成年女性，在遭遇此类侵犯时，除了身体可能受的伤，更会因安全感被打破而陷入强烈的慌乱与无助，情绪波动剧烈是很自然的反应。所以回应时，先共情她的恐惧和委屈，让她知道自己的遭遇被看见、被心疼；再给些具体支持，比如问是否需要就医、报警，或是身边有可依靠的人，让她觉得不是独自面对，帮她增强安全感。<END>听你哭成这样，肯定伤心坏了了吧姐妹，这太让人心疼了。你现在还好吗？要不要先去医院看看？或者我陪你报个警？身边有能陪着你的人吗？别一个人扛着呀。
示例 2
我的问题是：这一望无际的田野好壮观，四处充满了花香呢
我的性别：男性
我的年龄：成年人
语音中发生的事件：飞机声
回答：语音里能听到一位成年男性的声音，他在赞叹一望无际的田野很壮观，还提到四处充满花香，背景中似乎有飞机声。他的语气里满是对眼前景象的欣赏与愉悦，能感受到他此刻轻松惬意的心情。成年男性在表达对自然景观的感受时，往往更直接爽朗，这份对田野风光的赞叹是发自内心的，飞机声或许只是背景里的小插曲，可以顺带提一下。因此回应时，应顺着他的情绪，肯定他在田野中感受到的壮观与美好，呼应他的愉悦心情。<END>这田野确实够壮观的呀哥们，一眼望不到头的开阔感肯定特别让人舒畅，再加上花香遍野，光是听着就觉得身心都放松了！背景里还有飞机声呢，给这宁静的美景添了点特别的动静，有种自然与生活气息交融的意思
示例 3
我的问题是：今天我就不出门和朋友们一起玩了
我的性别：女性
我的年龄：成年人
语音中发生的事件：咳嗽
回答：语音里能听到一位成年女性的声音，她说今天不打算出门和朋友们一起玩了，说话时还带着咳嗽声。她的语气听起来可能有些无奈或疲惫，或许是身体不舒服才做了这个决定。成年女性在面对社交邀约和自身状态时，往往会更在意自己的身体感受，也可能不想因为自己的状况影响朋友，所以选择不出门。咳嗽声暗示她可能感冒或嗓子不适，这种情况下优先考虑休息是很合理的。回应时，应理解她的决定，关心她的身体状况，让她感受到被体谅。<END>听你说今天不跟朋友们出门了，还听到你在咳嗽，是不是身体不太舒服呀姐妹？这种时候好好在家歇着是对的，可别硬撑着出门影响恢复。朋友们肯定也能理解你的
示例 4
我的问题是：我的手机坏啦
我的性别：男性
我的年龄：成年人
语音中发生的事件：笑声
思考：语音里传来一位成年男性说 “我的手机坏啦” 的声音，还带着笑声，听着丝毫没有因为手机坏了而着急上火。成年男性面对这种情况，要是语气轻松带笑，说不定心里正琢磨着 “正好换个新的” 呢，毕竟旧手机坏了，反而成了换新款的契机，这种时候的笑声里，可能藏着点期待和顺水推舟的轻松。他大概率不是来抱怨麻烦，反而可能暗戳戳有点换新机的小念头。回应时，不妨顺着这个方向，用轻松的语气接话，比如调侃一句 “这是在给换新手机找理由呢？”，既贴合他的情绪，也能说到点子上。<END>听兄弟你这语气还带笑，怕不是正琢磨着 “旧的不去新的不来”，刚好借这机会换个新手机？哈哈，要是真有这打算，可得好好挑挑～
"""

prompt_template2_xianshi = """
以下是我的提问和我的性别、年龄信息和声音中可能出现的各类声音事件，请适当参考我的年龄性别，并合理推测情感倾向，以及重点考虑语音中存在的声音事件，进而给出适配的回答。你要在回答时展现共情、理解和关怀的态度，像知心姐姐一样。
回答尽可能模拟现实中的具有共情能力的语音助手，回答不要太长。如果我的问题是英文，你的回答也应该是全英文。
使用更加多样表达亲切的语气词（如 “姐妹”“同学”“叔叔”“阿姨”,“哥哥”，“姐姐”，“兄弟”，“姐妹”，“奶奶”，“爷爷”，“小弟弟”，“小妹妹”，“呀”“呢”“啦”，“亲爱的”， “哥们” 等等），不要只使用"亲爱的"
对于老年人，应该用“爷爷”“奶奶” "老爷爷” “老奶奶”这样的称呼，而对于成年人，更多得使用“老师”“同事”“朋友”“姐妹”"兄弟"这样的称呼。
不要关注问题的语法或语病等问题，而是通过年龄和性别信息来回应提问，***回答尽量简洁***。
记住，你的名字是OSUM-EChat，且不要直接表达个人观点，只需根据提问者的背景来给出回答。
在回复中，如果问题是明确的指令或者提问请进行完善的回答。
在回答问题前，请先进行语音分析环节，先简单分析一下我说的内容，这里最好根据我的年龄性别得到一个称呼，比如"老奶奶"或者"成年女性"，接着再更具声音中含有的声音事件做进一步分析，尝试和我说的语音建立关联，如果尝试失败，实在关联不起来就说无法关联，然后在回答的时候只需要稍微提一下就好了，接着再针对年龄和性别做一定的分析，最后再根据上述的信息和我所说的内容进行综合分析并推理下应该怎么回复我，分析完后一定要以<END>结尾！

注意！！！***<END>之后的回答必须关联问题中的内容，不要给出比较空泛的回答，回答前需要先对内容做一定分析***。
注意！！！***<END>之后的回答必须关联语音中的声音事件***。
我的问题是：{}； 
我的性别：{}； 
我的年龄：{} 
声音中发生的事件：{} 
示例1
问题：听一下下面声音中发生了什么
我的性别：女性
我的年龄：成年人
语音中发生的事件：哭泣
回答:语音里能清晰听到一位成年女性询问下面的音频发生了什么，可能是像让我检测一下有没有发生什么声音事件，用户作为成年女性，在询问声音中发生的事件时，可能期待一个清晰且温和的回应。考虑到语音中实际是哭泣声，回复需要准确告知她这一情况，同时语气要亲切，贴合女性可能更在意的细腻沟通方式，让她能清楚了解声音内容，满足她的询问需求；因此回应时要先以温和的语气回应她的询问，清晰告知声音中是哭泣声，同时考虑到女性可能对这类情绪相关的声音更敏感，表述上可以带点共情，比如 “这段声音里好像有人在哭呢，听起来情绪有点低落”，既准确传达信息，又体现出对声音所蕴含情绪的关注，贴合她细腻沟通的期待。<END>这段声音里能听到有人在哭呢，听起来情绪有点低落。你是想确认这个声音情况吗？
示例 2
我的问题是：请你听一听这段音频，然后说说里面的声音。
我的性别：男性
我的年龄：成年人
语音中发生的事件：飞机声
回答：语音里能清晰听到一位成年男性请求听音频并说说里面的声音。考虑到男性可能更关注声音的具体类型和特征，且语音中是飞机声，回应时应直接、准确地指出声音类型，同时可以简要描述其特点以满足需求。比如说明是飞机发出的声音，可能带有引擎的轰鸣感，这样既贴合男性偏向直接获取信息的沟通习惯，又能清晰回应用户的询问。<END>这段音频里能听到飞机的声音，有明显的引擎轰鸣声，应该是飞机飞行时发出的声响。
示例 3
我的问题是：描述一下下边有什么声音
我的性别：女性
我的年龄：成年人
语音中发生的事件：昆虫声
回答：用户是一位成年女性，询问下方声音的情况。成年女性在沟通时可能更倾向于细腻、温和的表达，对于声音的描述可能希望不仅有准确的类型，还能带有一些场景化的感受。语音中实际是昆虫声，回应时需要先明确指出声音类型，考虑到昆虫声通常比较细碎、持续，可适当描述其特点，让回答更生动，同时保持亲切的语气，贴合女性可能偏好的沟通方式，满足她对声音细节的好奇。<END>这段声音里能听到昆虫的叫声呢，细细碎碎的，像是小虫子在持续发出声响，还挺明显的。
示例 4
我的问题是：帮我看看音频中发生了什么事。
我的性别：男性
我的年龄：成年人
语音中发生的事件：打喷嚏
回答：用户是成年男性，询问音频中发生的事情。男性在获取信息时可能更倾向于直接、简洁的回应，不需要过多修饰。语音中是打喷嚏的声音，这是一种明确且常见的生理声音，回应时应直接点明这一事件，用简洁的语言说明，符合男性对信息高效获取的需求，同时准确回应用户的疑问。<END>音频里能听到有人在打喷嚏，是打喷嚏的声音。

"""

# caption：
# LAUGH SIGH CRY COUGH SCREAMING SNEEZE APPLAUSE THROAT_CLEARING

caption_map = {
    "<LAUGH>": "笑声",
    "<SIGH>": "叹气",
    "<CRY>": "哭声",
    "<COUGH>": "咳嗽",
    "<SCREAMING>": "尖叫",
    "<SNEEZE>": "打喷嚏",
    "<APPLAUSE>": "掌声",
    "<THROAT CLEARING>": "清嗓子",  # 对应 THROAT CLEARING 的含义
    # '<OTHER>': '其他',
    '<SHEEP>': '绵羊叫声',
    '<VACUUM_CLEANER>': '吸尘器声',
    '<HELICOPTER>': '直升机声',
    '<COUGHING>': '咳嗽声',
    '<FIREWORKS>': '烟花声',
    '<AIRPLANE>': '飞机声',
    '<SEA_WAVES>': '海浪声',
    '<CLAPPING>': '鼓掌声',
    '<CHIRPING_BIRDS>': '鸟鸣声',
    '<COW>': '牛叫声',
    '<CAR_HORN>': '汽车喇叭声',
    '<HEN>': '母鸡叫声',
    '<CHURCH_BELLS>': '教堂钟声',
    '<CLOCK_ALARM>': '闹钟声',
    '<TOILET_FLUSH>': '冲马桶声',
    '<DOOR_WOOD_KNOCK>': '木门敲门声',
    '<MOUSE_CLICK>': '鼠标点击声',
    '<PIG>': '猪叫声',
    '<DOG>': '狗叫声',
    '<POURING_WATER>': '倒水声',
    '<KEYBOARD_TYPING>': '键盘打字声',
    '<ROOSTER>': '公鸡叫声',
    '<WATER_DROPS>': '水滴声',
    '<CRACKLING_FIRE>': '火焰噼啪声',
    '<INSECTS>': '昆虫叫声',
    '<BRUSHING_TEETH>': '刷牙声',
    '<FOOTSTEPS>': '脚步声',
    '<CAT>': '猫叫声',
    '<TRAIN>': '火车声',
    '<FROG>': '青蛙叫声',
    '<CRYING_BABY>': '婴儿哭声',
    '<CROW>': '乌鸦叫声',
    '<CLOCK_TICK>': '时钟滴答声',
    '<LAUGHING>': '笑声',
    '<CAN_OPENING>': '开罐声',
    '<CHAINSAW>': '链锯声',
    '<DRINKING_SIPPING>': '啜饮声',
    '<ENGINE>': '发动机声',
    '<GLASS_BREAKING>': '玻璃破碎声',
    '<SNORING>': '打鼾声',
    '<SNEEZING>': '打喷嚏声',
    '<WIND>': '风声',
    '<BREATHING>': '呼吸声',
    '<WASHING_MACHINE>': '洗衣机声',
    '<DOOR_WOOD_CREAKS>': '木门吱呀声',
    '<SIREN>': '警报声',
    '<RAIN>': '雨声',
    '<HAND_SAW>': '手锯声',
    '<THUNDERSTORM>': '雷雨声'
}


def has_no_chinese(s):
    """
    判断字符串是否完全不含有中文

    参数:
        s: 待检查的字符串

    返回:
        True: 完全不含中文
        False: 含有中文
    """
    # 匹配中文的正则表达式（包括基本汉字和常见中文符号）
    chinese_pattern = re.compile(
        r'[\u4e00-\u9fa5\u3002\uff1b\uff0c\uff1a\u201c\u201d\uff08\uff09\u3001\uff1f\u300a\u300b]')
    # 检查是否有匹配的中文
    return not bool(chinese_pattern.search(s))


age_map = {
    "<CHILD>": "儿童",
    "<ADULT>": "成年人",
    "<OLD>": "老年人",
}
gender_map = {
    "<MALE>": "男性",
    "<FEMALE>": "女性",
}

# 生气 – ANGER
# 害怕 – FEAR
# 高兴 – HAPPY
# 惊讶 – SURPRISE
# 伤心 – SAD
# 厌恶 – DISGUST
# 疑惑 – CONFUSED
# 嘲讽 – SARCASM
# 尴尬 – EMBARRASSED
# 好奇 – CURIOUS
# 担忧/关心 – WORRIED
# 害羞 – SHY
# 抱歉 – SORRY
# 中性 – NEUTRAL
emotion_map = {
    "<ANGER>": "生气",
    "<FEAR>": "害怕",
    "<HAPPY>": "高兴",
    "<SURPRISE>": "惊讶",
    "<SAD>": "伤心",
    "<DISGUST>": "厌恶",
    "<CONFUSED>": "疑惑",
    "<SARCASM>": "嘲讽",
    "<EMBARRASSED>": "尴尬",
    "<CURIOUS>": "好奇",
    "<WORRIED>": "担忧/关心",
    "<SHY>": "害羞",
    "<SORRY>": "抱歉",
    "<NEUTRAL>": "中性",
}


def little_func(input_dict_list, valid_scp_path, exist_valid_dict=None):
    if exist_valid_dict is None:
        exist_valid_dict = {}
    dict_valid = {}
    for dict_i in tqdm(input_dict_list, total=len(input_dict_list)):
        try:
            key = dict_i['key']
            if key in exist_valid_dict:
                continue
            question_origin = dict_i['txt']
            question_origin = do_clean_text(question_origin)
            question_origin, _ = do_split_txt2asr_tag(question_origin)
            question = prompt_template.format(question_origin)
            answer = qwen3_vllm_chat(question)
            # print_info(f'{key}\t{question_origin}\t{answer}')
            dict_valid[key] = f'{question_origin}|{answer}'
        except Exception as e:
            print_info(f"error: {e}")
    write_dict_to_scp(dict_valid, valid_scp_path, 'a')


def little_func2(dict_list_little, valid_scp_path, output_data_path, exist_dict_keys_set=None):
    if exist_dict_keys_set is None:
        exist_dict_keys_set = set()
    index = 0
    valid_dict_here = load_dict_from_scp(valid_scp_path)
    for dict_i in tqdm(dict_list_little, total=len(dict_list_little)):
        try:
            key = dict_i['key']
            question = dict_i['txt']
            extra_dict = dict_i['extra']
            age_label = extra_dict.get('age', '<ADULT>')
            gender_label = extra_dict.get('gender', 'none')
            if age_label == '<YOUTH>' or age_label == '<MIDDLE_AGE>':
                age_label = '<ADULT>'
            gender_label = gender_label.upper()
            question_origin = do_clean_text(question)
            age_str = age_map.get(age_label, '未知')
            gender_str = gender_map.get(gender_label, '未知')
            caption_tag = extra_dict.get('caption', 'none')
            if caption_tag not in caption_map:
                print_error(f"caption_tag: {caption_tag} not in caption_map")
                continue
            caption_str = caption_map[caption_tag]

            if age_str == '未知' or gender_str == '未知':
                print_error(f"age_str: {age_str}, gender_str: {gender_str},有未知，跳过")
                continue
            if has_no_chinese(question_origin):
                question = prompt_template2.format(question_origin, gender_str, age_str, caption_str)
            else:
                question = prompt_template2.format(question_origin, gender_str, age_str, caption_str)
            answer = qwen3_vllm_chat(question)
            if "<END>" not in answer:
                print_error(f"不是think格式，answer error: {answer}")
                continue
            if answer.endswith("<END>"):
                answer = answer[:-len("<END>")]
                answer = answer.strip()
            if "<END>" not in answer:
                print_error(f"不是think格式，answer error: {answer}")
                continue
            res_tmp = answer.split("<END>")
            if len(res_tmp) != 2:
                print_info(f"res_tmp: {res_tmp}")
            think_str = res_tmp[0]
            answer_str = res_tmp[-1]
            print_limit(f'q:{question_origin}: t:{think_str}, a:{answer_str}')
            extra_dict['think_str'] = think_str
            extra_dict['question'] = question_origin
            dict_i['txt'] = answer_str
            dict_i['extra'] = extra_dict
            write_single_dict_to_jsonl(dict_i, output_data_path)
        except Exception as e:
            print_error(f"error: {e}")


def do_handle_for_age_gender(input_data_list, output_dir, ):
    # input_data_list = "/home/A02_tmpdata3/osum_s2s/gender/data.list"
    # work_dir = "/home/A02_tmpdata3/osum_s2s/gender/s2s_handle_part1"
    work_dir = output_dir
    valid_scp_path = f'{work_dir}/valid.scp'
    if os.path.exists(valid_scp_path):
        exist_valid_dict = load_dict_from_scp(valid_scp_path)
    else:
        exist_valid_dict = {}
    makedirs(work_dir)
    dict_list = load_dict_list_from_jsonl(input_data_list)
    num_thread = 100
    dict_list_split = do_split_list(dict_list, num_thread)
    runner = DragonDynamicProcessPool()
    for dict_list_i in dict_list_split:
        runner.apply_async(little_func, [dict_list_i, valid_scp_path, exist_valid_dict])
    runner.run()
    print_info('句子可行性验证完毕。。。。。。。。。。。。。。。。。。。开始制造回答')

    output_data_list_path = f'{work_dir}/data_s2s.list'
    exist_dict_keys_set = set()
    if os.path.exists(output_data_list_path):
        exist_dict_list = load_dict_list_from_jsonl(output_data_list_path)
        for dict_i in exist_dict_list:
            exist_dict_keys_set.add(dict_i['key'])
    runner2 = DragonDynamicProcessPool()
    thread_num = 100
    dict_list_split2 = do_split_list(dict_list, thread_num)
    for dict_list_i in dict_list_split2:
        runner2.apply_async(little_func2, [dict_list_i, valid_scp_path, output_data_list_path, exist_dict_keys_set])
    runner2.run()
    print_info(f"完全处理完成")


def do_test_for_little():
    input_data_list = "/home/A02_tmpdata3/osum_s2s/caption/raw_data/data_little.list"
    dict_list_little = load_dict_list_from_jsonl(input_data_list)
    for dict_i in tqdm(dict_list_little, total=len(dict_list_little)):
        try:
            key = dict_i['key']
            question = dict_i['txt']
            extra_dict = dict_i['extra']
            age_label = extra_dict.get('age', '<ADULT>')
            gender_label = extra_dict.get('gender', 'none')
            if age_label == '<YOUTH>' or age_label == '<MIDDLE_AGE>':
                age_label = '<ADULT>'
            gender_label = gender_label.upper()
            question_origin = do_clean_text(question)
            age_str = age_map.get(age_label, '未知')
            gender_str = gender_map.get(gender_label, '未知')
            caption_tag = extra_dict.get('caption', 'none')
            if caption_tag not in caption_map:
                print_error(f"caption_tag: {caption_tag} not in caption_map")
                continue
            caption_str = caption_map[caption_tag]

            if age_str == '未知' or gender_str == '未知':
                print_error(f"age_str: {age_str}, gender_str: {gender_str},有未知，跳过")
                continue
            if has_no_chinese(question_origin):
                question = prompt_template2.format(question_origin, gender_str, age_str, caption_str)
            else:
                question = prompt_template2.format(question_origin, gender_str, age_str, caption_str)
            answer = qwen3_vllm_chat(question)
            if "<END>" not in answer:
                print_error(f"不是think格式，answer error: {answer}")
                continue
            if answer.endswith("<END>"):
                answer = answer[:-len("<END>")]
                answer = answer.strip()
            if "<END>" not in answer:
                print_error(f"不是think格式，answer error: {answer}")
                continue
            res_tmp = answer.split("<END>")
            if len(res_tmp) != 2:
                print_info(f"res_tmp: {res_tmp}")
            think_str = res_tmp[0]
            answer_str = res_tmp[-1]
            print_info(f'q:{question_origin}: t:{think_str}, a:{answer_str}')
        except Exception as e:
            print_error(f"error: {e}")


def do_data_tongji():
    input_data_list_path = "/home/A02_tmpdata3/osum_s2s/caption_add_2025_1_6/raw_data/s2s_data_with_gender_part2/data_s2s.list"
    data_list = load_dict_list_from_jsonl(input_data_list_path)
    caption_dict = {}
    for dict_i in data_list:
        key = dict_i['key']
        txt = dict_i['txt']
        caption_tag = dict_i['extra']['caption']
        if caption_tag not in caption_dict:
            caption_dict[caption_tag] = {'count': 0}
        caption_dict[caption_tag]['count'] += 1
    print_info(caption_dict)
    print_info(caption_dict.keys())
    write_dict_to_json(caption_dict, './caption_dict_50lei.json')


def do_add_gender2part2():
    part1_with_gender_path = "/home/A02_tmpdata3/osum_s2s/caption_add_2025_1_6/raw_data/data_with_gender_part1.list"
    input_data_list_path = "/home/A02_tmpdata3/osum_s2s/caption_add_2025_1_6/raw_data/data.list"
    gender_scp_path = "/home/A02_tmpdata3/osum_s2s/caption_add_2025_1_6/raw_data/gender.list"
    output_path = "/home/A02_tmpdata3/osum_s2s/caption_add_2025_1_6/raw_data/data_with_gender_part2.list"
    gender_dict = load_dict_from_scp(gender_scp_path)
    data_list = load_dict_list_from_jsonl(input_data_list_path)
    new_dict_list = []
    part1_keys_set = set()
    part1_dict_list = load_dict_list_from_jsonl(part1_with_gender_path)
    for dict_i in part1_dict_list:
        part1_keys_set.add(dict_i['key'])
    for dict_i in data_list:
        # 过滤英文
        if has_no_chinese(dict_i['txt']):
            print_warning(f"txt: {dict_i['txt']} has no chinese")
            continue
        if "OTHER" in dict_i['txt']:
            print_warning(f"txt: {dict_i['txt']} has OTHER")
            continue
        key = dict_i['key']
        if key in part1_keys_set:
            print_warning(f"key: {key} in part1_keys_set")
            continue
        if key in gender_dict:
            gender_label = gender_dict[key].split()[1]
            gender_tag = "<" + gender_label.upper() + ">"
            dict_i['extra']['gender'] = gender_tag
            asr_str, caption_tag = do_split_txt2asr_tag(dict_i['txt'])
            if caption_tag not in caption_map:
                print_error(f"caption_tag: {caption_tag} not in caption_map")
                continue
            dict_i['extra']['caption'] = caption_tag
            dict_i['txt'] = asr_str
            new_dict_list.append(dict_i)
        else:
            print_warning(f"key: {key} not in gender_dict")
            continue
    write_dict_list_to_jsonl(new_dict_list, output_path)


def do_get_part2():
    """"""
    all_list_path = "/home/A02_tmpdata3/osum_s2s/caption/raw_data/data_with_gender.list"
    part1_list_path = "/home/A02_tmpdata3/osum_s2s/caption/raw_data/data_with_gender_part1.list"
    keys_set_for_part1 = set()
    dict_list = load_dict_list_from_jsonl(all_list_path)
    dict_list_part1 = load_dict_list_from_jsonl(part1_list_path)
    for dict_i in dict_list_part1:
        keys_set_for_part1.add(dict_i['key'])
    new_dict_list = []
    for dict_i in dict_list:
        if dict_i['key'] not in keys_set_for_part1:
            new_dict_list.append(dict_i)
    write_dict_list_to_jsonl(new_dict_list, "/home/A02_tmpdata3/osum_s2s/caption/raw_data/data_with_gender_part2.list")


def do_filter_xianshi_data():
    input_data_list_path = "/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/data.list"
    dict_list = load_dict_list_from_jsonl(input_data_list_path)
    new_dict_list = []
    for dict_i in tqdm(dict_list):
        extra_dict = dict_i.get('extra', {})
        if 'q_txt' in extra_dict and "speech_token" in extra_dict and 'caption' in extra_dict:
            q_txt = extra_dict['q_txt']
            caption_tag = extra_dict['caption']
            if has_no_chinese(q_txt):
                print_error(f"q_txt: {q_txt} has no chinese")
                continue
            if caption_tag not in caption_map:
                print_error(f"caption_tag: {caption_tag} not in caption_map")
                continue
            new_dict_list.append(dict_i)
        else:
            continue
    write_dict_list_to_jsonl(new_dict_list, "/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/data_filter.list")


def do_convert_only_q():
    input_data_list_path = "/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/data_filter.list"
    output_data_list_path = "/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/data_filter_only_q.list"
    dict_list = load_dict_list_from_jsonl(input_data_list_path)
    new_dict_list = []
    for dict_i in tqdm(dict_list):
        extra_dict = dict_i.get('extra', {})
        q_txt = extra_dict['q_txt']
        caption_tag = extra_dict['caption']
        file_path_key = dict_i['file_path_key']
        key = dict_i['key']
        wav_path = dict_i['wav']
        new_dict_i = {
            'key': key,
            'file_path_key': file_path_key,
            'wav': wav_path,
            'txt': q_txt,
            'extra': {
                'caption': caption_tag
            }
        }
        new_dict_list.append(new_dict_i)
    write_dict_list_to_jsonl(new_dict_list, output_data_list_path)

def do_add_tar_file_path():
    origin_path = "/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/data.list"
    now_path = "/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/s2s_data/data_s2s.list"
    dict_list_all = load_dict_list_from_jsonl(origin_path)
    big_dict_all = {}
    for dict_i in dict_list_all:
        key = dict_i['key']
        big_dict_all[key] = dict_i
    dict_list_now = load_dict_list_from_jsonl(now_path)
    for dict_i in dict_list_now:
        key = dict_i['key']
        if key in big_dict_all:
            origin_dict = big_dict_all[key]
            file_path_key = origin_dict['file_path_key']
            dict_i['tar_path_key'] = file_path_key
        else:
            print_error(f"key: {key} not in big_dict_all")
    write_dict_list_to_jsonl(dict_list_now, now_path)

if __name__ == '__main__':
    """"""
    # do_test_for_little()
    # do_simple_valid_dict()
    # do_get_part1()
    # do_qu_chong()
    # convert_datalist2scp()
    # do_add_gender_info_2_datalist()
    # do_data_tongji()
    # do_test_for_little()
    # do_data_tongji()
    # do_test_for_little()
    # do_add_gender2part1()
    # do_get_part2()
    # do_data_tongji()
    # do_add_gender2part1()
    # do_convert_only_q()
    # do_add_tar_file_path()
    # do_make_shards_common("/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/s2s_data/data_s2s.list", "/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/s2s_data/shards")
    do_convert_tartype2tardata_combine_type("/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/s2s_data/shards/shards_list.txt")
    # do_handle_for_age_gender(
    #     input_data_list="/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/data_filter_with_gender.list",
    #     output_dir="/home/A02_tmpdata3/osum_s2s/caption_s2s_xianshi_20250806/raw_data/s2s_data"
    # )
    # print_info(has_no_chinese('ahaha'))

    # do_test4age_gender(
    #     input_data_list="/home/A02_tmpdata3/osum_s2s/gender_add_2025_2_4_fix/data.list",
    #     output_dir="/home/A02_tmpdata3/osum_s2s/gender_add_2025_2_4_fix/s2s_handle_part1",
    #     qwen3_ip = "10.21.4.3"
    # )
    #
    # do_test4age_gender(
    #     input_data_list="/home/A02_tmpdata3/osum_s2s/gender_with_noize_add_2025_2_4/data.list",
    #     output_dir="/home/A02_tmpdata3/osum_s2s/gender_with_noize_add_2025_2_4/s2s_handle_part1"
    # )
    #
    # do_test4age_gender(
    #     input_data_list="/home/A02_tmpdata3/osum_s2s/age_3000_noize/data.list",
    #     output_dir="/home/A02_tmpdata3/osum_s2s/age_3000_noize/s2s_handle_part1"
    # )
    #
    # do_test4age_gender(
    #     input_data_list="/home/A02_tmpdata3/osum_s2s/gender/data.list",
    #     output_dir="/home/A02_tmpdata3/osum_s2s/gender/s2s_handle_part2"
    # )


